﻿import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
#from matplotlib.animation import FuncAnimation
#from matplotlib.animation import PillowWriter
from matplotlib.colors import ListedColormap
from matplotlib.colors import LinearSegmentedColormap
from numpy.core.umath import arccos
from matplotlib.patches import Arc
from matplotlib.animation import FFMpegWriter


# 设置参数
x_min, x_max = -20, 20  # x轴的范围
y_min, y_max = -20, 20  # y轴的范围
x = np.linspace(x_min, x_max, 1000)  # x轴的采样点
y = np.linspace(y_min, y_max, 1000)  # y轴的采样点
X, Y = np.meshgrid(x, y)  # 生成网格点

# 条件参数
time = 2
f = 2
wave_len = 2
n1 = 1.3
n2 = 1.6

# 中心亮斑的位置
source = (10, 0)


# 波的振幅函数
def wave_amplitude(x, y, t, source):  # （x,y)点在t时间由source引起的振幅
    r = 0.5 * np.sqrt((x - source[0]) ** 2 + (y - source[1]) ** 2)  # （x, y)与source的距离
    i = np.arctan(r / f)
    h = abs((time - t) * 5)  # 加绝对值还是会突变
    angle = (4 * np.pi / wave_len) * (2 * h) * np.sqrt((n1 ** 2 - (n2 * np.sin(i)) ** 2)) + np.pi
    return np.cos(angle / 2)


# 动画更新函数
def update_with_new_cmap(t):
    # 清除之前的图形
    plt.clf()
    Z = wave_amplitude(X, Y, t, source)
    plt.title("Newton's rings (time: {:.2f})".format(t))

    # 定义颜色映射的颜色
    colors = ['black', '#008000', 'green']  # 使用十六进制表示红色的中间色
    # 定义颜色映射的位置
    positions = [0.0, 0.5, 1.0]
    # 创建自定义的颜色映射
    cmap = LinearSegmentedColormap.from_list('CustomRedBlack', list(zip(positions, colors)))
    plt.imshow(Z, cmap=cmap, extent=[x_min, x_max, y_min, y_max], origin='lower')

    

    # 添加平行四边形表示高度
    basex = -15  # 平行四边形左下角坐标
    bottom_height = y_min + 4
    x= [-15,-3]
    
    top_height = bottom_height + abs((time - t) * 2)# 两平面间距
    thick=abs((time - t) * 5)
    plt.plot(x, [bottom_height,bottom_height], color='red', linewidth='1.0', linestyle='-')
    plt.plot(x,[top_height,top_height],color='steelblue')
    if thick>4.8:
        plt.plot([-15,-15],[bottom_height+1.9,top_height],color='steelblue')
        plt.plot([-3,-3],[bottom_height+1.9,top_height],color='steelblue')
        top=Arc((-9,bottom_height+1.9),12,3.7,angle=0,theta1=180,theta2=360,color='steelblue')
        plt.gca().add_patch(top)
    elif thick>2:
        top=Arc((-9,top_height),12,thick-0.5,angle=0,theta1=180,theta2=360,color='steelblue')
        plt.gca().add_patch(top)
    else:
        top=Arc((-9,top_height),12,thick,angle=0,theta1=180,theta2=360,color='steelblue')
        plt.gca().add_patch(top)
    # 显示高度值
    _y = y_min + 4 - 1

    plt.text(-14, _y, 'thickness: {:.2f}'.format(abs(time - t) * 5), color='b', fontsize=12)
    plt.text(basex-2, bottom_height, 'M1'.format(abs(time - t) * 5), color='b', fontsize=12)
    #plt.text(basex-1, top_height, 'M2'.format(abs(time - t) * 5), color='red', fontsize=12)
    
    #画出入射光线及反射光线
    plt.annotate("",
                xy=(-11, bottom_height),
                xytext=(-11, 8),
                # xycoords="figure points",
                arrowprops=dict(arrowstyle="->", color="r"))
    
    plt.annotate("",
                xy=(-7, bottom_height),
                xytext=(-7, 8),
                # xycoords="figure points",
                arrowprops=dict(arrowstyle="->", color="r"))
    
    plt.annotate("",
                xy=(-9, 8),
                xytext=(-9, bottom_height),
                # xycoords="figure points",
                arrowprops=dict(arrowstyle="->", color="g"))
    
    plt.annotate("",
                xy=(-13, 8),
                xytext=(-13, bottom_height),
                # xycoords="figure points",
                arrowprops=dict(arrowstyle="->", color="g"))
    
    plt.annotate("",
                xy=(-5, 8),
                xytext=(-5, bottom_height),
                # xycoords="figure points",
                arrowprops=dict(arrowstyle="->", color="g"))

    return plt

# 创建动画
fig = plt.figure(figsize=(8, 8))
frames = np.arange(0, 2 * time, 0.005)

ani = animation.FuncAnimation(fig, update_with_new_cmap, frames=frames, interval=150)

# 导出视频
writer = FFMpegWriter(fps=10, metadata=dict(artist='cm313'), bitrate=200000)
ani.save("Newton's rings.mp4", writer=writer)

# 显示动画
plt.show()
print("successfully")
